{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "self_attn_pool.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "HxUtAD83iXXM",
        "uMVXEMXHmMcA",
        "6mMFB9HwmjQr",
        "FYxVK6Joot3b",
        "5qI94Tc3o0e3",
        "v1LbdCBEpM_j",
        "eYLIUNfHpW99"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S-0trrYkwLp6",
        "colab_type": "text"
      },
      "source": [
        "# 第8章：图分类实战"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h43cBqYq0XrU",
        "colab_type": "text"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/FighterLYL/GraphNeuralNetwork/blob/master/chapter8/self_attn_pool.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HxUtAD83iXXM",
        "colab_type": "text"
      },
      "source": [
        "## SetUp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oLW4h0yW9a5K",
        "colab_type": "text"
      },
      "source": [
        "**注意**：`torch_scatter` 安装时编译需要用到cuda, 在colab中运行时可以选择`代码执行程序->更改运行时类型`,选择使用`GPU`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_QG_HXLziWWn",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "285b2611-f5e8-4b61-b56e-34093441de27"
      },
      "source": [
        "!pip install --verbose --no-cache-dir torch-scatter"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Created temporary directory: /tmp/pip-ephem-wheel-cache-l2yt13l1\n",
            "Created temporary directory: /tmp/pip-req-tracker-rv7k1nb6\n",
            "Created requirements tracker '/tmp/pip-req-tracker-rv7k1nb6'\n",
            "Created temporary directory: /tmp/pip-install-wpims5og\n",
            "Requirement already satisfied: torch-scatter in /usr/local/lib/python3.6/dist-packages (1.4.0)\n",
            "Cleaning up...\n",
            "Removed build tracker '/tmp/pip-req-tracker-rv7k1nb6'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CRHi8vpymFci",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import urllib\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.init as init\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data as data\n",
        "import torch.optim as optim\n",
        "import numpy as np\n",
        "import scipy.sparse as sp\n",
        "from zipfile import ZipFile\n",
        "from sklearn.model_selection import train_test_split\n",
        "import pickle\n",
        "import pandas as pd\n",
        "import torch_scatter\n",
        "from collections import Counter\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uMVXEMXHmMcA",
        "colab_type": "text"
      },
      "source": [
        "## 功能函数定义"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rP-9UIQhmUFD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def tensor_from_numpy(x, device):\n",
        "    return torch.from_numpy(x).to(device)\n",
        "\n",
        "\n",
        "def normalization(adjacency):\n",
        "    \"\"\"计算 L=D^-0.5 * (A+I) * D^-0.5,\n",
        "\n",
        "    Args:\n",
        "        adjacency: sp.csr_matrix.\n",
        "\n",
        "    Returns:\n",
        "        归一化后的邻接矩阵，类型为 torch.sparse.FloatTensor\n",
        "    \"\"\"\n",
        "    adjacency += sp.eye(adjacency.shape[0])    # 增加自连接\n",
        "    degree = np.array(adjacency.sum(1))\n",
        "    d_hat = sp.diags(np.power(degree, -0.5).flatten())\n",
        "    L = d_hat.dot(adjacency).dot(d_hat).tocoo()\n",
        "    # 转换为 torch.sparse.FloatTensor\n",
        "    indices = torch.from_numpy(np.asarray([L.row, L.col])).long()\n",
        "    values = torch.from_numpy(L.data.astype(np.float32))\n",
        "    tensor_adjacency = torch.sparse.FloatTensor(indices, values, L.shape)\n",
        "    return tensor_adjacency\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6mMFB9HwmjQr",
        "colab_type": "text"
      },
      "source": [
        "## D&D数据"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-VeHbOJxHYy",
        "colab_type": "text"
      },
      "source": [
        "D&D is a dataset of 1178 protein structures (Dobson and Doig, 2003). Each protein is represented by a graph, in which the nodes are amino acids and two nodes are connected by an edge if they are less than 6 Angstroms apart. The prediction task is to classify the protein structures into enzymes and non-enzymes.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iAfYjmh-mnz2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class DDDataset(object):\n",
        "    url = \"https://ls11-www.cs.tu-dortmund.de/people/morris/graphkerneldatasets/DD.zip\"\n",
        "    \n",
        "    def __init__(self, data_root=\"data\", train_size=0.8):\n",
        "        self.data_root = data_root\n",
        "        self.maybe_download()\n",
        "        sparse_adjacency, node_labels, graph_indicator, graph_labels = self.read_data()\n",
        "        self.sparse_adjacency = sparse_adjacency.tocsr()\n",
        "        self.node_labels = node_labels\n",
        "        self.graph_indicator = graph_indicator\n",
        "        self.graph_labels = graph_labels\n",
        "        self.train_index, self.test_index = self.split_data(train_size)\n",
        "        self.train_label = graph_labels[self.train_index]\n",
        "        self.test_label = graph_labels[self.test_index]\n",
        "\n",
        "    def split_data(self, train_size):\n",
        "        unique_indicator = np.asarray(list(set(self.graph_indicator)))\n",
        "        train_index, test_index = train_test_split(unique_indicator,\n",
        "                                                   train_size=train_size,\n",
        "                                                   random_state=1234)\n",
        "        return train_index, test_index\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        mask = self.graph_indicator == index\n",
        "        node_labels = self.node_labels[mask]\n",
        "        graph_indicator = self.graph_indicator[mask]\n",
        "        graph_labels = self.graph_labels[index]\n",
        "        adjacency = self.sparse_adjacency[mask, :][:, mask]\n",
        "        return adjacency, node_labels, graph_indicator, graph_labels\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.graph_labels)\n",
        "    \n",
        "    def read_data(self):\n",
        "        data_dir = os.path.join(self.data_root, \"DD\")\n",
        "        print(\"Loading DD_A.txt\")\n",
        "        adjacency_list = np.genfromtxt(os.path.join(data_dir, \"DD_A.txt\"),\n",
        "                                       dtype=np.int64, delimiter=',') - 1\n",
        "        print(\"Loading DD_node_labels.txt\")\n",
        "        node_labels = np.genfromtxt(os.path.join(data_dir, \"DD_node_labels.txt\"), \n",
        "                                    dtype=np.int64) - 1\n",
        "        print(\"Loading DD_graph_indicator.txt\")\n",
        "        graph_indicator = np.genfromtxt(os.path.join(data_dir, \"DD_graph_indicator.txt\"), \n",
        "                                        dtype=np.int64) - 1\n",
        "        print(\"Loading DD_graph_labels.txt\")\n",
        "        graph_labels = np.genfromtxt(os.path.join(data_dir, \"DD_graph_labels.txt\"), \n",
        "                                     dtype=np.int64) - 1\n",
        "        num_nodes = len(node_labels)\n",
        "        sparse_adjacency = sp.coo_matrix((np.ones(len(adjacency_list)), \n",
        "                                          (adjacency_list[:, 0], adjacency_list[:, 1])),\n",
        "                                         shape=(num_nodes, num_nodes), dtype=np.float32)\n",
        "        print(\"Number of nodes: \", num_nodes)\n",
        "        return sparse_adjacency, node_labels, graph_indicator, graph_labels\n",
        "    \n",
        "    def maybe_download(self):\n",
        "        save_path = os.path.join(self.data_root)\n",
        "        if not os.path.exists(save_path):\n",
        "            self.download_data(self.url, save_path)\n",
        "        if not os.path.exists(os.path.join(self.data_root, \"DD\")):\n",
        "            zipfilename = os.path.join(self.data_root, \"DD.zip\")\n",
        "            with ZipFile(zipfilename, \"r\") as zipobj:\n",
        "                zipobj.extractall(os.path.join(self.data_root))\n",
        "                print(\"Extracting data from {}\".format(zipfilename))\n",
        "    \n",
        "    @staticmethod\n",
        "    def download_data(url, save_path):\n",
        "        \"\"\"数据下载工具，当原始数据不存在时将会进行下载\"\"\"\n",
        "        print(\"Downloading data from {}\".format(url))\n",
        "        if not os.path.exists(save_path):\n",
        "            os.makedirs(save_path)\n",
        "        data = urllib.request.urlopen(url)\n",
        "        filename = \"DD.zip\"\n",
        "        with open(os.path.join(save_path, filename), 'wb') as f:\n",
        "            f.write(data.read())\n",
        "        return True\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1kFvRNKmrv0",
        "colab_type": "text"
      },
      "source": [
        "## Model定义"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYxVK6Joot3b",
        "colab_type": "text"
      },
      "source": [
        "### GraphConvolution"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGKsuL5toxB8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class GraphConvolution(nn.Module):\n",
        "    def __init__(self, input_dim, output_dim, use_bias=True):\n",
        "        \"\"\"图卷积：L*X*\\theta\n",
        "\n",
        "        Args:\n",
        "        ----------\n",
        "            input_dim: int\n",
        "                节点输入特征的维度\n",
        "            output_dim: int\n",
        "                输出特征维度\n",
        "            use_bias : bool, optional\n",
        "                是否使用偏置\n",
        "        \"\"\"\n",
        "        super(GraphConvolution, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.output_dim = output_dim\n",
        "        self.use_bias = use_bias\n",
        "        self.weight = nn.Parameter(torch.Tensor(input_dim, output_dim))\n",
        "        if self.use_bias:\n",
        "            self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
        "        else:\n",
        "            self.register_parameter('bias', None)\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        init.kaiming_uniform_(self.weight)\n",
        "        if self.use_bias:\n",
        "            init.zeros_(self.bias)\n",
        "\n",
        "    def forward(self, adjacency, input_feature):\n",
        "        \"\"\"邻接矩阵是稀疏矩阵，因此在计算时使用稀疏矩阵乘法\"\"\"\n",
        "        support = torch.mm(input_feature, self.weight)\n",
        "        output = torch.sparse.mm(adjacency, support)\n",
        "        if self.use_bias:\n",
        "            output += self.bias\n",
        "        return output\n",
        "\n",
        "    def __repr__(self):\n",
        "        return self.__class__.__name__ + ' (' \\\n",
        "            + str(self.input_dim) + ' -> ' \\\n",
        "            + str(self.output_dim) + ')'\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5qI94Tc3o0e3",
        "colab_type": "text"
      },
      "source": [
        "### ReadOut 实现"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZM21IkMo93V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def global_max_pool(x, graph_indicator):\n",
        "    num = graph_indicator.max().item() + 1\n",
        "    return torch_scatter.scatter_max(x, graph_indicator, dim=0, dim_size=num)[0]\n",
        "\n",
        "\n",
        "def global_avg_pool(x, graph_indicator):\n",
        "    num = graph_indicator.max().item() + 1\n",
        "    return torch_scatter.scatter_mean(x, graph_indicator, dim=0, dim_size=num)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v1LbdCBEpM_j",
        "colab_type": "text"
      },
      "source": [
        "### 基于自注意力机制的池化层"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xh14Np1WqBEa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def top_rank(attention_score, graph_indicator, keep_ratio):\n",
        "    \"\"\"基于给定的attention_score, 对每个图进行pooling操作.\n",
        "    为了直观体现pooling过程，我们将每个图单独进行池化，最后再将它们级联起来进行下一步计算\n",
        "    \n",
        "    Arguments:\n",
        "    ----------\n",
        "        attention_score：torch.Tensor\n",
        "            使用GCN计算出的注意力分数，Z = GCN(A, X)\n",
        "        graph_indicator：torch.Tensor\n",
        "            指示每个节点属于哪个图\n",
        "        keep_ratio: float\n",
        "            要保留的节点比例，保留的节点数量为int(N * keep_ratio)\n",
        "    \"\"\"\n",
        "    # TODO: 确认是否是有序的, 必须是有序的\n",
        "    graph_id_list = list(set(graph_indicator.cpu().numpy()))\n",
        "    mask = attention_score.new_empty((0,), dtype=torch.bool)\n",
        "    for graph_id in graph_id_list:\n",
        "        graph_attn_score = attention_score[graph_indicator == graph_id]\n",
        "        graph_node_num = len(graph_attn_score)\n",
        "        graph_mask = attention_score.new_zeros((graph_node_num,),\n",
        "                                                dtype=torch.bool)\n",
        "        keep_graph_node_num = int(keep_ratio * graph_node_num)\n",
        "        _, sorted_index = graph_attn_score.sort(descending=True)\n",
        "        graph_mask[sorted_index[:keep_graph_node_num]] = True\n",
        "        mask = torch.cat((mask, graph_mask))\n",
        "    \n",
        "    return mask\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jrLzPXHhp6HB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def filter_adjacency(adjacency, mask):\n",
        "    \"\"\"根据掩码mask对图结构进行更新\n",
        "    \n",
        "    Args:\n",
        "        adjacency: torch.sparse.FloatTensor, 池化之前的邻接矩阵\n",
        "        mask: torch.Tensor(dtype=torch.bool), 节点掩码向量\n",
        "    \n",
        "    Returns:\n",
        "        torch.sparse.FloatTensor, 池化之后归一化邻接矩阵\n",
        "    \"\"\"\n",
        "    device = adjacency.device\n",
        "    mask = mask.cpu().numpy()\n",
        "    indices = adjacency.coalesce().indices().cpu().numpy()\n",
        "    num_nodes = adjacency.size(0)\n",
        "    row, col = indices\n",
        "    maskout_self_loop = row != col\n",
        "    row = row[maskout_self_loop]\n",
        "    col = col[maskout_self_loop]\n",
        "    sparse_adjacency = sp.csr_matrix((np.ones(len(row)), (row, col)),\n",
        "                                     shape=(num_nodes, num_nodes), dtype=np.float32)\n",
        "    filtered_adjacency = sparse_adjacency[mask, :][:, mask]\n",
        "    return normalization(filtered_adjacency).to(device)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gj_f0-XnuNnA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SelfAttentionPooling(nn.Module):\n",
        "    def __init__(self, input_dim, keep_ratio, activation=torch.tanh):\n",
        "        super(SelfAttentionPooling, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.keep_ratio = keep_ratio\n",
        "        self.activation = activation\n",
        "        self.attn_gcn = GraphConvolution(input_dim, 1)\n",
        "    \n",
        "    def forward(self, adjacency, input_feature, graph_indicator):\n",
        "        attn_score = self.attn_gcn(adjacency, input_feature).squeeze()\n",
        "        attn_score = self.activation(attn_score)\n",
        "        \n",
        "        mask = top_rank(attn_score, graph_indicator, self.keep_ratio)\n",
        "        hidden = input_feature[mask] * attn_score[mask].view(-1, 1)\n",
        "        mask_graph_indicator = graph_indicator[mask]\n",
        "        mask_adjacency = filter_adjacency(adjacency, mask)\n",
        "        return hidden, mask_graph_indicator, mask_adjacency\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eYLIUNfHpW99",
        "colab_type": "text"
      },
      "source": [
        "### 模型一：SAGPool Global Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rMHS7xicpUMj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ModelA(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, num_classes=2):\n",
        "        \"\"\"图分类模型结构A\n",
        "        \n",
        "        Args:\n",
        "        ----\n",
        "            input_dim: int, 输入特征的维度\n",
        "            hidden_dim: int, 隐藏层单元数\n",
        "            num_classes: 分类类别数 (default: 2)\n",
        "        \"\"\"\n",
        "        super(ModelA, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.num_classes = num_classes\n",
        "        \n",
        "        self.gcn1 = GraphConvolution(input_dim, hidden_dim)\n",
        "        self.gcn2 = GraphConvolution(hidden_dim, hidden_dim)\n",
        "        self.gcn3 = GraphConvolution(hidden_dim, hidden_dim)\n",
        "        self.pool = SelfAttentionPooling(hidden_dim * 3, 0.5)\n",
        "        self.fc1 = nn.Linear(hidden_dim * 3 * 2, hidden_dim)\n",
        "        self.fc2 = nn.Linear(hidden_dim, hidden_dim // 2)\n",
        "        self.fc3 = nn.Linear(hidden_dim // 2, num_classes)\n",
        "\n",
        "    def forward(self, adjacency, input_feature, graph_indicator):\n",
        "        gcn1 = F.relu(self.gcn1(adjacency, input_feature))\n",
        "        gcn2 = F.relu(self.gcn2(adjacency, gcn1))\n",
        "        gcn3 = F.relu(self.gcn3(adjacency, gcn2))\n",
        "        \n",
        "        gcn_feature = torch.cat((gcn1, gcn2, gcn3), dim=1)\n",
        "        pool, pool_graph_indicator, pool_adjacency = self.pool(adjacency, gcn_feature,\n",
        "                                                               graph_indicator)\n",
        "        \n",
        "        readout = torch.cat((global_avg_pool(pool, pool_graph_indicator),\n",
        "                             global_max_pool(pool, pool_graph_indicator)), dim=1)\n",
        "        \n",
        "        fc1 = F.relu(self.fc1(readout))\n",
        "        fc2 = F.relu(self.fc2(fc1))\n",
        "        logits = self.fc3(fc2)\n",
        "        \n",
        "        return logits\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FqXt-VdqpcsC",
        "colab_type": "text"
      },
      "source": [
        "### 模型二：SAGPool Hierarchical Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zx2wCs0xppUN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ModelB(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, num_classes=2):\n",
        "        \"\"\"图分类模型结构\n",
        "        \n",
        "        Args:\n",
        "        -----\n",
        "            input_dim: int, 输入特征的维度\n",
        "            hidden_dim: int, 隐藏层单元数\n",
        "            num_classes: int, 分类类别数 (default: 2)\n",
        "        \"\"\"\n",
        "        super(ModelB, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.num_classes = num_classes\n",
        "        \n",
        "        self.gcn1 = GraphConvolution(input_dim, hidden_dim)\n",
        "        self.pool1 = SelfAttentionPooling(hidden_dim, 0.5)\n",
        "        self.gcn2 = GraphConvolution(hidden_dim, hidden_dim)\n",
        "        self.pool2 = SelfAttentionPooling(hidden_dim, 0.5)\n",
        "        self.gcn3 = GraphConvolution(hidden_dim, hidden_dim)\n",
        "        self.pool3 = SelfAttentionPooling(hidden_dim, 0.5)\n",
        "        \n",
        "        self.mlp = nn.Sequential(\n",
        "            nn.Linear(hidden_dim * 2, hidden_dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(hidden_dim, hidden_dim // 2),\n",
        "            nn.ReLU(), \n",
        "            nn.Linear(hidden_dim // 2, num_classes))\n",
        "    \n",
        "    def forward(self, adjacency, input_feature, graph_indicator):\n",
        "        gcn1 = F.relu(self.gcn1(adjacency, input_feature))\n",
        "        pool1, pool1_graph_indicator, pool1_adjacency = \\\n",
        "            self.pool1(adjacency, gcn1, graph_indicator)\n",
        "        global_pool1 = torch.cat(\n",
        "            [global_avg_pool(pool1, pool1_graph_indicator),\n",
        "             global_max_pool(pool1, pool1_graph_indicator)],\n",
        "            dim=1)\n",
        "        \n",
        "        gcn2 = F.relu(self.gcn2(pool1_adjacency, pool1))\n",
        "        pool2, pool2_graph_indicator, pool2_adjacency = \\\n",
        "            self.pool2(pool1_adjacency, gcn2, pool1_graph_indicator)\n",
        "        global_pool2 = torch.cat(\n",
        "            [global_avg_pool(pool2, pool2_graph_indicator),\n",
        "             global_max_pool(pool2, pool2_graph_indicator)],\n",
        "            dim=1)\n",
        "\n",
        "        gcn3 = F.relu(self.gcn3(pool2_adjacency, pool2))\n",
        "        pool3, pool3_graph_indicator, pool3_adjacency = \\\n",
        "            self.pool3(pool2_adjacency, gcn3, pool2_graph_indicator)\n",
        "        global_pool3 = torch.cat(\n",
        "            [global_avg_pool(pool3, pool3_graph_indicator),\n",
        "             global_max_pool(pool3, pool3_graph_indicator)],\n",
        "            dim=1)\n",
        "        \n",
        "        readout = global_pool1 + global_pool2 + global_pool3\n",
        "        \n",
        "        logits = self.mlp(readout)\n",
        "        return logits\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94kE4RkdmvFO",
        "colab_type": "text"
      },
      "source": [
        "## 训练 & 测试"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "unoZOcTqeFuI",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        },
        "outputId": "31966846-ddf9-46e2-d561-8ff951556a7e"
      },
      "source": [
        "dataset = DDDataset()"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading DD_A.txt\n",
            "Loading DD_node_labels.txt\n",
            "Loading DD_graph_indicator.txt\n",
            "Loading DD_graph_labels.txt\n",
            "Number of nodes:  334925\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XLXGbQvHiiMT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 模型输入数据准备\n",
        "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "adjacency = dataset.sparse_adjacency\n",
        "normalize_adjacency = normalization(adjacency).to(DEVICE)\n",
        "node_labels = tensor_from_numpy(dataset.node_labels, DEVICE)\n",
        "node_features = F.one_hot(node_labels, node_labels.max().item() + 1).float()\n",
        "graph_indicator = tensor_from_numpy(dataset.graph_indicator, DEVICE)\n",
        "graph_labels = tensor_from_numpy(dataset.graph_labels, DEVICE)\n",
        "train_index = tensor_from_numpy(dataset.train_index, DEVICE)\n",
        "test_index = tensor_from_numpy(dataset.test_index, DEVICE)\n",
        "train_label = tensor_from_numpy(dataset.train_label, DEVICE)\n",
        "test_label = tensor_from_numpy(dataset.test_label, DEVICE)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BQv_zcThATJk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 超参数设置\n",
        "INPUT_DIM = node_features.size(1)\n",
        "NUM_CLASSES = 2\n",
        "EPOCHS = 200    # @param {type: \"integer\"}\n",
        "HIDDEN_DIM =    32# @param {type: \"integer\"}\n",
        "LEARNING_RATE = 0.01 # @param\n",
        "WEIGHT_DECAY = 0.0001 # @param"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9xjOe4gg4hV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 模型初始化\n",
        "model_g = ModelA(INPUT_DIM, HIDDEN_DIM, NUM_CLASSES).to(DEVICE)\n",
        "model_h = ModelB(INPUT_DIM, HIDDEN_DIM, NUM_CLASSES).to(DEVICE)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B5cDpOhEtdcS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = model_g #@param ['model_g', 'model_h'] {type: 'raw'}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "myPjiGN-6qxk",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 221
        },
        "outputId": "2eb3d30a-a687-4e6d-94fa-7227ea3846b2"
      },
      "source": [
        "print(\"Device:\", DEVICE)\n",
        "print(model)"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Device: cuda\n",
            "ModelA(\n",
            "  (gcn1): GraphConvolution (89 -> 32)\n",
            "  (gcn2): GraphConvolution (32 -> 32)\n",
            "  (gcn3): GraphConvolution (32 -> 32)\n",
            "  (pool): SelfAttentionPooling(\n",
            "    (attn_gcn): GraphConvolution (96 -> 1)\n",
            "  )\n",
            "  (fc1): Linear(in_features=192, out_features=32, bias=True)\n",
            "  (fc2): Linear(in_features=32, out_features=16, bias=True)\n",
            "  (fc3): Linear(in_features=16, out_features=2, bias=True)\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gtGt1GXDsS1c",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "3c47b6f4-58b7-4fa6-ba27-9bb1a101726e"
      },
      "source": [
        "criterion = nn.CrossEntropyLoss().to(DEVICE)\n",
        "optimizer = optim.Adam(model.parameters(), LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n",
        "\n",
        "model.train()\n",
        "for epoch in range(EPOCHS):\n",
        "    logits = model(normalize_adjacency, node_features, graph_indicator)\n",
        "    loss = criterion(logits[train_index], train_label)  # 只对训练的数据计算损失值\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()  # 反向传播计算参数的梯度\n",
        "    optimizer.step()  # 使用优化方法进行梯度更新\n",
        "    train_acc = torch.eq(\n",
        "        logits[train_index].max(1)[1], train_label).float().mean()\n",
        "    print(\"Epoch {:03d}: Loss {:.4f}, TrainAcc {:.4}\".format(\n",
        "        epoch, loss.item(), train_acc.item()))\n"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 000: Loss 0.6814, TrainAcc 0.5881\n",
            "Epoch 001: Loss 0.6738, TrainAcc 0.5881\n",
            "Epoch 002: Loss 0.6696, TrainAcc 0.5881\n",
            "Epoch 003: Loss 0.6697, TrainAcc 0.5881\n",
            "Epoch 004: Loss 0.6579, TrainAcc 0.5881\n",
            "Epoch 005: Loss 0.6558, TrainAcc 0.5881\n",
            "Epoch 006: Loss 0.6556, TrainAcc 0.5881\n",
            "Epoch 007: Loss 0.6409, TrainAcc 0.621\n",
            "Epoch 008: Loss 0.6472, TrainAcc 0.6253\n",
            "Epoch 009: Loss 0.6333, TrainAcc 0.6486\n",
            "Epoch 010: Loss 0.6418, TrainAcc 0.6656\n",
            "Epoch 011: Loss 0.6265, TrainAcc 0.6794\n",
            "Epoch 012: Loss 0.6302, TrainAcc 0.6486\n",
            "Epoch 013: Loss 0.6256, TrainAcc 0.6561\n",
            "Epoch 014: Loss 0.6144, TrainAcc 0.6837\n",
            "Epoch 015: Loss 0.6190, TrainAcc 0.6773\n",
            "Epoch 016: Loss 0.6074, TrainAcc 0.673\n",
            "Epoch 017: Loss 0.6082, TrainAcc 0.6741\n",
            "Epoch 018: Loss 0.5945, TrainAcc 0.7028\n",
            "Epoch 019: Loss 0.5962, TrainAcc 0.6932\n",
            "Epoch 020: Loss 0.5945, TrainAcc 0.6815\n",
            "Epoch 021: Loss 0.5802, TrainAcc 0.6975\n",
            "Epoch 022: Loss 0.5957, TrainAcc 0.6911\n",
            "Epoch 023: Loss 0.5765, TrainAcc 0.7017\n",
            "Epoch 024: Loss 0.5723, TrainAcc 0.7059\n",
            "Epoch 025: Loss 0.5775, TrainAcc 0.7059\n",
            "Epoch 026: Loss 0.5564, TrainAcc 0.7208\n",
            "Epoch 027: Loss 0.5672, TrainAcc 0.7038\n",
            "Epoch 028: Loss 0.5503, TrainAcc 0.7314\n",
            "Epoch 029: Loss 0.5504, TrainAcc 0.7261\n",
            "Epoch 030: Loss 0.5501, TrainAcc 0.7325\n",
            "Epoch 031: Loss 0.5334, TrainAcc 0.7473\n",
            "Epoch 032: Loss 0.5397, TrainAcc 0.7378\n",
            "Epoch 033: Loss 0.5310, TrainAcc 0.7537\n",
            "Epoch 034: Loss 0.5167, TrainAcc 0.7633\n",
            "Epoch 035: Loss 0.5205, TrainAcc 0.7484\n",
            "Epoch 036: Loss 0.5152, TrainAcc 0.7601\n",
            "Epoch 037: Loss 0.5032, TrainAcc 0.7696\n",
            "Epoch 038: Loss 0.4942, TrainAcc 0.7749\n",
            "Epoch 039: Loss 0.4987, TrainAcc 0.7675\n",
            "Epoch 040: Loss 0.5036, TrainAcc 0.7654\n",
            "Epoch 041: Loss 0.4926, TrainAcc 0.7707\n",
            "Epoch 042: Loss 0.4745, TrainAcc 0.7792\n",
            "Epoch 043: Loss 0.4665, TrainAcc 0.7887\n",
            "Epoch 044: Loss 0.4720, TrainAcc 0.7909\n",
            "Epoch 045: Loss 0.4740, TrainAcc 0.7856\n",
            "Epoch 046: Loss 0.4643, TrainAcc 0.7866\n",
            "Epoch 047: Loss 0.4439, TrainAcc 0.8057\n",
            "Epoch 048: Loss 0.4410, TrainAcc 0.8174\n",
            "Epoch 049: Loss 0.4406, TrainAcc 0.8089\n",
            "Epoch 050: Loss 0.4481, TrainAcc 0.7962\n",
            "Epoch 051: Loss 0.4479, TrainAcc 0.7941\n",
            "Epoch 052: Loss 0.4427, TrainAcc 0.81\n",
            "Epoch 053: Loss 0.4199, TrainAcc 0.8185\n",
            "Epoch 054: Loss 0.4096, TrainAcc 0.828\n",
            "Epoch 055: Loss 0.4233, TrainAcc 0.8217\n",
            "Epoch 056: Loss 0.4124, TrainAcc 0.8174\n",
            "Epoch 057: Loss 0.4039, TrainAcc 0.8174\n",
            "Epoch 058: Loss 0.3880, TrainAcc 0.8418\n",
            "Epoch 059: Loss 0.3856, TrainAcc 0.8397\n",
            "Epoch 060: Loss 0.3801, TrainAcc 0.845\n",
            "Epoch 061: Loss 0.3951, TrainAcc 0.8259\n",
            "Epoch 062: Loss 0.4363, TrainAcc 0.8079\n",
            "Epoch 063: Loss 0.4916, TrainAcc 0.7781\n",
            "Epoch 064: Loss 0.3780, TrainAcc 0.8376\n",
            "Epoch 065: Loss 0.4109, TrainAcc 0.8132\n",
            "Epoch 066: Loss 0.4094, TrainAcc 0.8185\n",
            "Epoch 067: Loss 0.3742, TrainAcc 0.845\n",
            "Epoch 068: Loss 0.4071, TrainAcc 0.8227\n",
            "Epoch 069: Loss 0.3514, TrainAcc 0.8599\n",
            "Epoch 070: Loss 0.3744, TrainAcc 0.8408\n",
            "Epoch 071: Loss 0.3902, TrainAcc 0.8301\n",
            "Epoch 072: Loss 0.3523, TrainAcc 0.8631\n",
            "Epoch 073: Loss 0.3513, TrainAcc 0.8631\n",
            "Epoch 074: Loss 0.3575, TrainAcc 0.8482\n",
            "Epoch 075: Loss 0.3306, TrainAcc 0.8684\n",
            "Epoch 076: Loss 0.3474, TrainAcc 0.8652\n",
            "Epoch 077: Loss 0.3293, TrainAcc 0.8694\n",
            "Epoch 078: Loss 0.3246, TrainAcc 0.8684\n",
            "Epoch 079: Loss 0.3383, TrainAcc 0.8588\n",
            "Epoch 080: Loss 0.3142, TrainAcc 0.8715\n",
            "Epoch 081: Loss 0.3091, TrainAcc 0.8715\n",
            "Epoch 082: Loss 0.3208, TrainAcc 0.8652\n",
            "Epoch 083: Loss 0.3008, TrainAcc 0.8843\n",
            "Epoch 084: Loss 0.2949, TrainAcc 0.8843\n",
            "Epoch 085: Loss 0.2961, TrainAcc 0.8811\n",
            "Epoch 086: Loss 0.3088, TrainAcc 0.8694\n",
            "Epoch 087: Loss 0.3171, TrainAcc 0.8694\n",
            "Epoch 088: Loss 0.3067, TrainAcc 0.8779\n",
            "Epoch 089: Loss 0.2829, TrainAcc 0.8896\n",
            "Epoch 090: Loss 0.2742, TrainAcc 0.8938\n",
            "Epoch 091: Loss 0.2979, TrainAcc 0.879\n",
            "Epoch 092: Loss 0.2934, TrainAcc 0.8811\n",
            "Epoch 093: Loss 0.2984, TrainAcc 0.8705\n",
            "Epoch 094: Loss 0.2862, TrainAcc 0.8875\n",
            "Epoch 095: Loss 0.2663, TrainAcc 0.9002\n",
            "Epoch 096: Loss 0.2744, TrainAcc 0.8907\n",
            "Epoch 097: Loss 0.2755, TrainAcc 0.8864\n",
            "Epoch 098: Loss 0.2942, TrainAcc 0.8737\n",
            "Epoch 099: Loss 0.2602, TrainAcc 0.8949\n",
            "Epoch 100: Loss 0.2623, TrainAcc 0.8949\n",
            "Epoch 101: Loss 0.2310, TrainAcc 0.9087\n",
            "Epoch 102: Loss 0.2521, TrainAcc 0.9023\n",
            "Epoch 103: Loss 0.2297, TrainAcc 0.9193\n",
            "Epoch 104: Loss 0.2344, TrainAcc 0.9066\n",
            "Epoch 105: Loss 0.2192, TrainAcc 0.9214\n",
            "Epoch 106: Loss 0.2253, TrainAcc 0.9151\n",
            "Epoch 107: Loss 0.2289, TrainAcc 0.913\n",
            "Epoch 108: Loss 0.2159, TrainAcc 0.9246\n",
            "Epoch 109: Loss 0.2179, TrainAcc 0.9161\n",
            "Epoch 110: Loss 0.1981, TrainAcc 0.9352\n",
            "Epoch 111: Loss 0.2073, TrainAcc 0.9225\n",
            "Epoch 112: Loss 0.1861, TrainAcc 0.9374\n",
            "Epoch 113: Loss 0.1945, TrainAcc 0.9289\n",
            "Epoch 114: Loss 0.1860, TrainAcc 0.9384\n",
            "Epoch 115: Loss 0.1747, TrainAcc 0.9448\n",
            "Epoch 116: Loss 0.1841, TrainAcc 0.9374\n",
            "Epoch 117: Loss 0.1655, TrainAcc 0.9427\n",
            "Epoch 118: Loss 0.1717, TrainAcc 0.9374\n",
            "Epoch 119: Loss 0.1635, TrainAcc 0.948\n",
            "Epoch 120: Loss 0.1607, TrainAcc 0.9522\n",
            "Epoch 121: Loss 0.2336, TrainAcc 0.8854\n",
            "Epoch 122: Loss 0.7919, TrainAcc 0.6805\n",
            "Epoch 123: Loss 1.0422, TrainAcc 0.6996\n",
            "Epoch 124: Loss 0.5043, TrainAcc 0.7537\n",
            "Epoch 125: Loss 0.9418, TrainAcc 0.5732\n",
            "Epoch 126: Loss 0.6531, TrainAcc 0.6805\n",
            "Epoch 127: Loss 0.5462, TrainAcc 0.7622\n",
            "Epoch 128: Loss 0.6121, TrainAcc 0.7537\n",
            "Epoch 129: Loss 0.5723, TrainAcc 0.7665\n",
            "Epoch 130: Loss 0.3502, TrainAcc 0.8641\n",
            "Epoch 131: Loss 0.4173, TrainAcc 0.8068\n",
            "Epoch 132: Loss 0.4610, TrainAcc 0.7845\n",
            "Epoch 133: Loss 0.3979, TrainAcc 0.8376\n",
            "Epoch 134: Loss 0.4896, TrainAcc 0.8068\n",
            "Epoch 135: Loss 0.3175, TrainAcc 0.8662\n",
            "Epoch 136: Loss 0.3619, TrainAcc 0.8333\n",
            "Epoch 137: Loss 0.3281, TrainAcc 0.8652\n",
            "Epoch 138: Loss 0.3243, TrainAcc 0.8684\n",
            "Epoch 139: Loss 0.3672, TrainAcc 0.8355\n",
            "Epoch 140: Loss 0.3262, TrainAcc 0.8641\n",
            "Epoch 141: Loss 0.3300, TrainAcc 0.8747\n",
            "Epoch 142: Loss 0.3460, TrainAcc 0.862\n",
            "Epoch 143: Loss 0.3255, TrainAcc 0.8662\n",
            "Epoch 144: Loss 0.3167, TrainAcc 0.8673\n",
            "Epoch 145: Loss 0.3247, TrainAcc 0.8567\n",
            "Epoch 146: Loss 0.3124, TrainAcc 0.8641\n",
            "Epoch 147: Loss 0.2965, TrainAcc 0.897\n",
            "Epoch 148: Loss 0.2926, TrainAcc 0.9045\n",
            "Epoch 149: Loss 0.2841, TrainAcc 0.9013\n",
            "Epoch 150: Loss 0.2726, TrainAcc 0.8949\n",
            "Epoch 151: Loss 0.2704, TrainAcc 0.8981\n",
            "Epoch 152: Loss 0.2707, TrainAcc 0.8949\n",
            "Epoch 153: Loss 0.2539, TrainAcc 0.914\n",
            "Epoch 154: Loss 0.2610, TrainAcc 0.9087\n",
            "Epoch 155: Loss 0.2481, TrainAcc 0.914\n",
            "Epoch 156: Loss 0.2373, TrainAcc 0.9119\n",
            "Epoch 157: Loss 0.2391, TrainAcc 0.9066\n",
            "Epoch 158: Loss 0.2221, TrainAcc 0.9268\n",
            "Epoch 159: Loss 0.2219, TrainAcc 0.9204\n",
            "Epoch 160: Loss 0.2163, TrainAcc 0.9299\n",
            "Epoch 161: Loss 0.2099, TrainAcc 0.9321\n",
            "Epoch 162: Loss 0.2086, TrainAcc 0.9299\n",
            "Epoch 163: Loss 0.2014, TrainAcc 0.9363\n",
            "Epoch 164: Loss 0.1983, TrainAcc 0.9395\n",
            "Epoch 165: Loss 0.1936, TrainAcc 0.9416\n",
            "Epoch 166: Loss 0.1899, TrainAcc 0.9384\n",
            "Epoch 167: Loss 0.1855, TrainAcc 0.9395\n",
            "Epoch 168: Loss 0.1820, TrainAcc 0.9416\n",
            "Epoch 169: Loss 0.1776, TrainAcc 0.9459\n",
            "Epoch 170: Loss 0.1739, TrainAcc 0.9469\n",
            "Epoch 171: Loss 0.1704, TrainAcc 0.948\n",
            "Epoch 172: Loss 0.1662, TrainAcc 0.9501\n",
            "Epoch 173: Loss 0.1641, TrainAcc 0.9512\n",
            "Epoch 174: Loss 0.1599, TrainAcc 0.9522\n",
            "Epoch 175: Loss 0.1572, TrainAcc 0.9522\n",
            "Epoch 176: Loss 0.1532, TrainAcc 0.9554\n",
            "Epoch 177: Loss 0.1505, TrainAcc 0.9575\n",
            "Epoch 178: Loss 0.1472, TrainAcc 0.9586\n",
            "Epoch 179: Loss 0.1442, TrainAcc 0.9575\n",
            "Epoch 180: Loss 0.1408, TrainAcc 0.9628\n",
            "Epoch 181: Loss 0.1383, TrainAcc 0.966\n",
            "Epoch 182: Loss 0.1349, TrainAcc 0.965\n",
            "Epoch 183: Loss 0.1325, TrainAcc 0.9671\n",
            "Epoch 184: Loss 0.1292, TrainAcc 0.9682\n",
            "Epoch 185: Loss 0.1263, TrainAcc 0.9713\n",
            "Epoch 186: Loss 0.1237, TrainAcc 0.9724\n",
            "Epoch 187: Loss 0.1206, TrainAcc 0.9713\n",
            "Epoch 188: Loss 0.1176, TrainAcc 0.9724\n",
            "Epoch 189: Loss 0.1151, TrainAcc 0.9724\n",
            "Epoch 190: Loss 0.1127, TrainAcc 0.9724\n",
            "Epoch 191: Loss 0.1096, TrainAcc 0.9735\n",
            "Epoch 192: Loss 0.1074, TrainAcc 0.9735\n",
            "Epoch 193: Loss 0.1048, TrainAcc 0.9735\n",
            "Epoch 194: Loss 0.1021, TrainAcc 0.9745\n",
            "Epoch 195: Loss 0.0996, TrainAcc 0.9745\n",
            "Epoch 196: Loss 0.0973, TrainAcc 0.9756\n",
            "Epoch 197: Loss 0.0946, TrainAcc 0.9756\n",
            "Epoch 198: Loss 0.0924, TrainAcc 0.9777\n",
            "Epoch 199: Loss 0.0906, TrainAcc 0.9798\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "foY0A1iAuFAR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "be9fa76a-51c2-448f-c2b9-a4576a83a398"
      },
      "source": [
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    logits = model(normalize_adjacency, node_features, graph_indicator)\n",
        "    test_logits = logits[test_index]\n",
        "    test_acc = torch.eq(\n",
        "        test_logits.max(1)[1], test_label\n",
        "    ).float().mean()\n",
        "\n",
        "print(test_acc.item())"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "0.7076271176338196\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ivo1wTiJu8Dm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}